//
// Sample: Blocking IPv4/IPv6 Client
//
// Files:
//      bclient.cpp     - this file
//      resolve.cpp     - routines for resovling addresses, etc.
//      resolve.h       - header file for resolve.c
//
// Description:
//      This sample illustrates simple blocking IO for TCP and UDP for
//      both IPv4 and IPv6. This sample uses the getaddrinfo/getnameinfo
//      APIs which allows this application to be IP agnostic. That is the
//      desired address family (AF_INET or AF_INET6) can be determined
//      simply from the string address passed via the -l and -n command.
//
// Compile:
//      cl -o bclient.exe bclient.cpp resolve.cpp ws2_32.lib
//
// Usage:
//      bclient.exe [options]
//          -a 4|6     Address family, 4 = IPv4, 6 = IPv6 [default = IPv4]\n"
//          -e port    Port number [default = 5150]\n"
//          -l addr    Local address to bind to [default INADDR_ANY for IPv4 or INADDR6_ANY for IPv6
//          -n addr    Remote address to connect/send to
//          -b size    Buffer size (in bytes)
//          -r addr    Relay address
//          -p port    Relay port
//          -x count   Number of sends to perform
//
#include <winsock2.h>
#include <ws2tcpip.h>
#include <stdio.h>
#include <stdlib.h>

#include "resolve.h"

#define DEFAULT_BUFFER_SIZE     4096
#define DEFAULT_SEND_COUNT      100

int gAddressFamily = AF_UNSPEC,             // Address family to use        
gBufferSize    = DEFAULT_BUFFER_SIZE,   // Default buffer size for sends
gSendCount     = DEFAULT_SEND_COUNT;    // Default number of sends to perform

char *gBindAddr    = NULL,                  // Address to bind to locally
*gServerAddr  = "127.0.0.1",                  // Server name/address
*gRelayAddr   = "127.0.0.1",                  // Relay address
*gRelayPort   = "8585",                // Relay port
*gServerPort    = "5150";              // Port to connect to (on server side)

char* REQUEST_CONNECT = "connect to 127.0.0.1:8585",
*RESPONSE_OK    = "connect ok",
*RESPONSE_FAIL = "connect fail";

struct addrinfo *gConnectedEndpoint=NULL;   // Addresses that server name/address resolved to
struct addrinfo *gRelayEndpoint=NULL;       // Addresses that server relay server wait

// Statistics variables
volatile  LONG gBytesRead=0,
gBytesSent=0,
gStartTime=0;

struct _BUFFER_OBJ;

//
// Allocated for each client connection
//
typedef struct _CONNECTION_OBJ
{
    SOCKET      s;              // Client socket
    HANDLE      hRecvSema;      // Semaphore incremented for each receive

    struct _BUFFER_OBJ *PendingSendHead,    // List of pending buffers to send
        *PendingSendTail;    // End of the list

    CRITICAL_SECTION    SendRecvQueueCritSec; // Protect access to this structure

} CONNECTION_OBJ;

//
// Allocate for bouncing two socket
//
typedef struct _BOUNCING_PARAM
{
    CONNECTION_OBJ* ConnObjFrom;
    CONNECTION_OBJ* ConnObjTo;
}BOUNCING_PARAM;

//
// Allocate for bouncing two socket
//
typedef struct _SERVER_PAIRE
{
    SOCKET serverSocket;
    SOCKET relaySocket;
}SERVER_PAIRE;

//
// Allocated for each receiver posted
//    Each receive thread allocates one of these for a receive operation.
//    After data is read, this object is queued for the send thread to
//    echo back to the client (sender).
//
typedef struct _BUFFER_OBJ
{
    char        *buf;           // Data buffer for data
    int          buflen;        // Length of buffer or number of bytes contained in buffer    
    struct _BUFFER_OBJ *next;

} BUFFER_OBJ;

CONNECTION_OBJ *GetConnectionObj(SOCKET s)
{
    CONNECTION_OBJ *newobj=NULL;

    // Allocate the object
    newobj = (CONNECTION_OBJ *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(CONNECTION_OBJ));
    if (newobj == NULL)
    {
        printf("GetConnectionObj: HeapAlloc failed: %d\n", WSAGetLastError());
        ExitProcess(-1);
    }
    newobj->s = s;

    // Create the semaphore for signaling the send thread
    newobj->hRecvSema = CreateSemaphore(NULL, 0, 0x0FFFFFFF, NULL);
    if (newobj->hRecvSema == NULL)
    {
        fprintf(stderr, "GetConnectionObj: CreateSemaphore failed: %d\n", GetLastError());
        ExitProcess(-1);
    }

    InitializeCriticalSection(&newobj->SendRecvQueueCritSec);

    return newobj;
}

//
// Function: FreeConnectionObj
//
// Description:
//    This routine frees a CONNECTIN_OBJ. It first frees the critical section.
//    See the comment for GetConnectionObj about using lookaside lists.
//
void FreeConnectionObj(CONNECTION_OBJ *obj)
{
    DeleteCriticalSection(&obj->SendRecvQueueCritSec);
    HeapFree(GetProcessHeap(), 0, obj);
}

//
// Function: GetBufferObj
// 
// Description:
//    Allocate a BUFFER_OBJ. Each receive posted by a receive thread allocates
//    one of these. After the recv is successful, the BUFFER_OBJ is queued for
//    sending by the send thread. Again, lookaside lists may be used to increase
//    performance.
//
BUFFER_OBJ *GetBufferObj(int buflen)
{
    BUFFER_OBJ *newobj=NULL;

    // Allocate the object
    newobj = (BUFFER_OBJ *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(BUFFER_OBJ));
    if (newobj == NULL)
    {
        fprintf(stderr, "GetBufferObj: HeapAlloc failed: %d\n", GetLastError());
        ExitProcess(-1);
    }
    // Allocate the buffer
    newobj->buf = (char *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(BYTE) *buflen);
    if (newobj->buf == NULL)
    {
        fprintf(stderr, "GetBufferObj: HeapAlloc failed: %d\n", GetLastError());
        ExitProcess(-1);
    }
    newobj->buflen = buflen;

    return newobj;
}

//
// Function: FreeBufferObj
// 
// Description:
//    Free the buffer object.
//
void FreeBufferObj(BUFFER_OBJ *obj)
{
    HeapFree(GetProcessHeap(), 0, obj->buf);
    HeapFree(GetProcessHeap(), 0, obj);
}

//
// Function: EnqueueBufferObj
//
// Description:
//   Queue up a receive buffer for this connection. After enqueueing the receiver
//   will release the counting semaphore which will release the sender to 
//   dequeue a buffer and send it.
//
void EnqueueBufferObj(CONNECTION_OBJ *conn, BUFFER_OBJ *obj)
{
    EnterCriticalSection(&conn->SendRecvQueueCritSec);

    if (conn->PendingSendHead == NULL)
    {
        // Queue is empty
        conn->PendingSendHead = conn->PendingSendTail = obj;
    }
    else
    {
        // Put new object at the end 
        conn->PendingSendTail->next = obj;
        conn->PendingSendTail = obj;
    }
    LeaveCriticalSection(&conn->SendRecvQueueCritSec);
}

// 
// Function: DequeueBufferObj
//
// Description:
//    Remove a BUFFER_OBJ from the given connection's queue for sending.
//
BUFFER_OBJ *DequeueBufferObj(CONNECTION_OBJ *conn)
{
    BUFFER_OBJ *ret=NULL;

    EnterCriticalSection(&conn->SendRecvQueueCritSec);

    if (conn->PendingSendTail != NULL)
    {
        // Queue is non empty
        ret = conn->PendingSendHead;

        conn->PendingSendHead = conn->PendingSendHead->next;
        if (conn->PendingSendTail == ret)
        {
            // Item is the only item in the queue
            conn->PendingSendTail = NULL;
        }
    }

    LeaveCriticalSection(&conn->SendRecvQueueCritSec);

    return ret;
}

DWORD WINAPI SendThread(LPVOID lpParam)
{
    CONNECTION_OBJ *ConnObj=NULL;
    BUFFER_OBJ     *BuffObj=NULL;
    int             rc,
        nleft,
        idx;


    // Retrieve the connection object
    ConnObj = (CONNECTION_OBJ *)lpParam;

    while (1)
    {
        // Wait for the receive thread to signal us
        rc = WaitForSingleObject(ConnObj->hRecvSema, INFINITE);
        if (rc == WAIT_FAILED || rc == WAIT_TIMEOUT)
        {
            fprintf(stderr, "WaitForSingleObject failed: %d\n", GetLastError());
            ExitProcess(-1);
        }

        // Retrieve the first buffer from this connection's queue
        BuffObj = DequeueBufferObj(ConnObj);

        //
        // If the this receive by the receive thread indicated zero bytes then
        // the connection has been gracefully closed. Otherwise, if an error
        // was indicated then the connection was aborted.
        //
        if (((BuffObj->buflen == 0) || (BuffObj->buflen == SOCKET_ERROR)))
        {
            FreeBufferObj(BuffObj);
            BuffObj = NULL;
            break;
        }



        // Otherwise send the buffer on the connection socket
        nleft = BuffObj->buflen;
        idx = 0;

        while (nleft > 0)
        {
            rc = send(
                ConnObj->s,
                &BuffObj->buf[idx],
                nleft,
                0
                );
            if (rc == SOCKET_ERROR)
            {
                break;
            }
            else
            {
                nleft -= rc;
                idx += rc;
            }
        }

        if (rc == SOCKET_ERROR)
        {
            printf("SendThread: send(to) failed: %d\n", WSAGetLastError());
            break;
        }
        else if (rc > 0)
        {
            // Increment the statistics
            InterlockedExchangeAdd(&gBytesSent, rc);            
        }

        FreeBufferObj(BuffObj);
        BuffObj = NULL;
    }

    // Close the connection's socket
    closesocket(ConnObj->s);

    FreeConnectionObj(ConnObj);

    ExitThread(0);
    return 0;
}

//
DWORD WINAPI BouncThread(LPVOID lpParam)
{
    BOUNCING_PARAM *BouncParam = NULL;
    CONNECTION_OBJ *ConnObjFrom=NULL;
    CONNECTION_OBJ *ConnObjTo=NULL;
    BUFFER_OBJ     *BuffObj=NULL;
    int             rc;

    // Retrieve the connection object
    BouncParam = (BOUNCING_PARAM *)lpParam;
    ConnObjFrom = BouncParam->ConnObjFrom;
    ConnObjTo = BouncParam->ConnObjTo;

    // loop until the connection is closed or is aborted/terminated
    while (1)
    {
        // Allocate the buffer for stream send/recv
        BuffObj = GetBufferObj(gBufferSize);

        rc = recv(
            ConnObjFrom->s, 
            BuffObj->buf, 
            BuffObj->buflen, 
            0);
        BuffObj->buflen = rc;

        // Queue the receive buffer for sending and signal the send thread
        EnqueueBufferObj(ConnObjTo, BuffObj);

        ReleaseSemaphore(ConnObjTo->hRecvSema, 1, NULL);

        if (rc == 0 || rc == SOCKET_ERROR)
        {
            break;
        }
        else if (rc != SOCKET_ERROR)
        {
            // Increment the statistics
            InterlockedExchangeAdd(&gBytesRead, rc);            
        }

    }

    ExitThread(0);
    return 0;
}

//
// Function: usage
//
// Description:
//      Prints usage information and exits the process.
//
void usage(char *progname)
{
    fprintf(stderr, "usage: %s [-a 4|6] [-e port] [-l local-addr] [-n addr] [-p udp|tcp]\n",
        progname);
    fprintf(stderr, 
        "  -a 4|6     Address family, 4 = IPv4, 6 = IPv6 [default = IPv4]\n"
        "  -e port    Port number [default = 5150]\n"
        "  -l addr    Local address to bind to [default INADDR_ANY for IPv4 or INADDR6_ANY for IPv6]\n"
        "  -n addr    Remote address to connect/send to\n"            
        "  -b size    Buffer size\n"
        "  -x count   Number of sends to perform\n"
        "  -r addr    Relay address"
        "  -p port    Relay port"
        );
    ExitProcess(-1);
}

//
// Function: ValidateArgs
//
// Description:
//      Parses the command line arguments and sets up some global 
//      variables.
//
void ValidateArgs(int argc, char **argv)
{
    int     i;

    for(i=1; i < argc ;i++)
    {
        if (((argv[i][0] != '/') && (argv[i][0] != '-')) || (strlen(argv[i]) < 2))
            usage(argv[0]);
        else
        {
            switch (tolower(argv[i][1]))
            {
            case 'a':               // address family - IPv4 or IPv6
                if (i+1 >= argc)
                    usage(argv[0]);
                if (argv[i+1][0] == '4')
                    gAddressFamily = AF_INET;
                else if (argv[i+1][0] == '6')
                    gAddressFamily = AF_INET6;
                else
                    usage(argv[0]);
                i++;
                break;
            case 'b':               // buffer size
                if (i+1 >= argc)
                    usage(argv[0]);
                gBufferSize = atol(argv[++i]);
                break;                
            case 'e':               // endpoint - port number
                if (i+1 >= argc)
                    usage(argv[0]);
                gServerPort = argv[++i];
                break;
            case 'l':               // local address for binding
                if (i+1 >= argc)
                    usage(argv[0]);
                gBindAddr = argv[++i];
                break;
            case 'n':               // address to connect/send to
                if (i+1 >= argc)
                    usage(argv[0]);
                gServerAddr = argv[++i];
                break;

            case 'r':               // relay server address
                if (i+1 >= argc)
                    usage(argv[0]);
                gRelayAddr = argv[++i];
                break;

            case 'p':               // relay server port
                if (i+1 >= argc)
                    usage(argv[0]);
                gRelayPort = argv[++i];
                break;

            case 'x':               // send count
                if (i+1 >=argc)
                    usage(argv[0]);
                gSendCount = atoi(argv[++i]);
                break;
            default:
                usage(argv[0]);
                break;
            }
        }
    }
}

//
// Function: main
//
// Description:
//      This is the main program. It parses the command line and creates
//      the main socket. For UDP this socket is used to receive datagrams.
//      For TCP the socket is used to accept incoming client connections.
//      Each client TCP connection is handed off to a worker thread which
//      will receive any data on that connection until the connection is
//      closed.
//
int __cdecl main(int argc, char **argv)
{
    WSADATA          wsd;
    SOCKET           s;
    int              rc;                            // return code    

    // Parse the command line
    ValidateArgs(argc, argv);

    // Load Winsock
    if (WSAStartup(MAKEWORD(2,2), &wsd) != 0)
    {
        fprintf(stderr, "unable to load Winsock!\n");
        return -1;
    }

    // Resolve the server's name
    gConnectedEndpoint = ResolveAddress(gServerAddr, gServerPort, gAddressFamily, SOCK_STREAM, IPPROTO_TCP);
    if (gConnectedEndpoint == NULL)
    {
        fprintf(stderr, "ResolveAddress failed to return any addresses!\n");
        return -1;
    }

    // Resolve the relay server's name
    gRelayEndpoint = ResolveAddress(gRelayAddr, gRelayPort, gAddressFamily, SOCK_STREAM, IPPROTO_TCP);
    if (gRelayEndpoint == NULL)
    {
        fprintf(stderr, "ResolveAddress failed to return any addresses!\n");
        return -1;
    }

    // initial done, start working
    while (TRUE)
    {        
        // create the socket
        s = socket(gConnectedEndpoint->ai_family, gConnectedEndpoint->ai_socktype, gConnectedEndpoint->ai_protocol);
        if (s == INVALID_SOCKET)
        {
            fprintf(stderr, "socket failed: %d\n", WSAGetLastError());
            return -1;
        }

        // connect to server
        do 
        {
            rc = connect(s, gConnectedEndpoint->ai_addr, gConnectedEndpoint->ai_addrlen);
            if (rc == SOCKET_ERROR)
            {
                printf("connect failed: %d\n", WSAGetLastError());
                printf("sleep 5' before continue\n");
                Sleep(5000);
                closesocket(s);
                s = INVALID_SOCKET;
            }else{
                printf("connected to server ");
                PrintAddress(gConnectedEndpoint->ai_addr,gConnectedEndpoint->ai_addrlen);
                printf("\n");
            }       
        } while (rc==SOCKET_ERROR);

        // process connect request
        bool bOneMoreRelay = false;
        char* buff = new char[256];

        while(!bOneMoreRelay){

            printf("receive connect request...\n");
            rc = recv(s,buff,256,NULL);
            if(rc == SOCKET_ERROR){
                fprintf(stderr, "receive failed: %d\n", WSAGetLastError());
                closesocket(s);
                s=INVALID_SOCKET;
                break;
            }
            buff[rc] = 0;
            printf("received request\n");   
            printf(buff);
            printf("\n");

            // create the bouncing socket
            SOCKET sbnc = socket(gRelayEndpoint->ai_family, gRelayEndpoint->ai_socktype, gRelayEndpoint->ai_protocol);
            if (sbnc == INVALID_SOCKET)
            {
                fprintf(stderr, "socket failed: %d\n", WSAGetLastError());

                //send back not valid message
                send(s,RESPONSE_FAIL,strlen(RESPONSE_FAIL),NULL);
                break;
            }
            
            rc = connect(sbnc, gRelayEndpoint->ai_addr, gRelayEndpoint->ai_addrlen);
            if (rc == SOCKET_ERROR)
            {
                printf("connect failed: %d\n", WSAGetLastError());
                printf("sleep 5' before continue\n");
                Sleep(5000);
                closesocket(sbnc);
                sbnc = INVALID_SOCKET;

                //send back not valid message
                send(s,RESPONSE_FAIL,strlen(RESPONSE_FAIL),NULL);
                break;
            }  


            // bouncing it
            send(s,RESPONSE_OK,strlen(RESPONSE_OK),NULL);

            printf("send back ok, create send thread\n");
            CONNECTION_OBJ* ConnObjFrom = GetConnectionObj(s);
            CONNECTION_OBJ* ConnObjTo = GetConnectionObj(sbnc);
            HANDLE hThread = CreateThread(NULL, 0, SendThread, (LPVOID)ConnObjFrom, 0, NULL);
            if (hThread == NULL)
            {
                fprintf(stderr, "CreateThread failed: %d\n", GetLastError());
                ExitThread(-1);
            }
            CloseHandle(hThread);

            hThread = CreateThread(NULL, 0, SendThread, (LPVOID)ConnObjTo, 0, NULL);
            if (hThread == NULL)
            {
                fprintf(stderr, "CreateThread failed: %d\n", GetLastError());
                ExitThread(-1);
            }
            CloseHandle(hThread);

            printf("create bouncing thread\n");
            BOUNCING_PARAM* pBnc = new BOUNCING_PARAM;
            pBnc->ConnObjFrom = ConnObjFrom;
            pBnc->ConnObjTo = ConnObjTo;

            hThread = CreateThread(NULL, 0, BouncThread, (LPVOID)pBnc, 0, NULL);
            if (hThread == NULL)
            {
                fprintf(stderr, "CreateThread failed: %d\n", GetLastError());
                ExitThread(-1);
            }
            CloseHandle(hThread);

            BOUNCING_PARAM* pBnc2 = new BOUNCING_PARAM;
            pBnc2->ConnObjFrom = ConnObjTo;
            pBnc2->ConnObjTo = ConnObjFrom;

            hThread = CreateThread(NULL, 0, BouncThread, (LPVOID)pBnc2, 0, NULL);
            if (hThread == NULL)
            {
                fprintf(stderr, "CreateThread failed: %d\n", GetLastError());
                ExitThread(-1);
            }
            CloseHandle(hThread);
            bOneMoreRelay = true;                        
        }        
    }   

    freeaddrinfo(gConnectedEndpoint);
    freeaddrinfo(gRelayEndpoint);

    WSACleanup();
    return 0;
}
